#pragma once

#include <vector>
#include <stdint.h>
#include <random>
#include <memory>
#include <queue>
#include "tensor_data.h"

namespace  ldl
{
namespace core
{
template<typename T>
class Function;

template<typename T>
class Exp;

template<typename T>
class Add;

template<typename T>
class Mul;

template<typename T>
class Tensor : public std::enable_shared_from_this<Tensor<T>>
{
public:
// 初始化列表构造函数（修改为更安全的实现）
    Tensor(std::initializer_list<T> init_list) : m_data(init_list) {
        m_grad = TensorData<T>(m_data.shape(), 0);
    }

    // 多维初始化列表构造函数（安全版）
    template<typename U>
    Tensor(std::initializer_list<std::initializer_list<U>> init_list) : m_data (init_list){
        m_grad = TensorData<T>(m_data.shape(), 0);
    }

    Tensor(const std::vector<int64_t> &shape, const std::vector<T> &data) : m_data(shape, data) {
        m_grad = TensorData<T>(m_data.shape(), 0);
    }

    Tensor(const std::vector<int64_t> &shape, const T& data) : m_data(shape, data) {
        m_grad = TensorData<T>(m_data.shape(), 0);
    }
    Tensor(const Tensor&& other)
    {
        m_data = other.m_data;
        m_grad = other.m_grad;
        m_creator = other.m_creator;
    }
    Tensor()
    {

    }
    friend std::ostream& operator<<(std::ostream& os, const Tensor<T>& tensor)
    {
        return os << tensor.m_data;
    }

    int64_t size()
    {
        return m_data.size();
    }

    /**
     * @brief 取非const左值的data，data也允许修改
     */
    TensorData<T>& data() &
    {
        return m_data;
    }
    TensorData<T>& grad() &
    {
        return m_grad;
    }

    /**
     * @brief 取const左值的data，因而返回的data也应该是const
     */
    const TensorData<T>& data() const &
    {
        return m_data;
    }
    const TensorData<T>& grad() const &
    {
        return m_grad;
    }

    const std::vector<int64_t>& shape() const
    {
        return m_data.shape();;
    }

    std::shared_ptr<Tensor<T>> exp() {
        std::shared_ptr<Function<T>> creator = std::make_shared<Exp<T>>();
        return (*creator)(this->shared_from_this());
    }

    /**
     * @brief 加法
     * 
     * @param other 
     * @return std::shared_ptr<Tensor<T>> 
     */
    std::shared_ptr<Tensor<T>> operator+(Tensor<T>& other) {
        std::shared_ptr<Function<T>> creator = std::make_shared<Add<T>>();
        return operate(other, creator);
    }

    /**
     * @brief 按元素相乘
     */
    std::shared_ptr<Tensor<T>> operator*(Tensor<T>& other) {
        std::shared_ptr<Function<T>> creator = std::make_shared<Mul<T>>();
        return operate(other, creator);
    }

    std::shared_ptr<Tensor<T>> operate(Tensor<T>& other, std::shared_ptr<Function<T>> creator) {
        std::vector<std::shared_ptr<Tensor<T>>> inputs{};
        inputs.push_back(this->shared_from_this());
        inputs.push_back(other.get_shared_ptr());
        return (*creator)(inputs)[0];
    }


    void backward() {
        m_grad = TensorData<T>(m_data.shape(), 1);
        std::queue<std::shared_ptr<Function<T>>> creators{};
        creators.push(m_creator);
        while(!creators.empty()) {
            auto creator = creators.front();
            creators.pop();
            creator->backward();
            for(auto item:creator->inputs()) {
                if(item->creator() != nullptr) {
                    creators.push(item->creator());
                }
            }
        } 
    }

    void set_creator(std::shared_ptr<Function<T>> creator)
    {
        m_creator = creator;
    }
    std::shared_ptr<Function<T>> creator() {
        return m_creator;
    }

    std::shared_ptr<Tensor<T>> get_shared_ptr() {
        return this->shared_from_this();
    }
private:
    TensorData<T> m_data{};
    TensorData<T> m_grad{};
    std::shared_ptr<Function<T>> m_creator{};
};
}
}
